import os
import requests
import json
import uuid
import logging
from typing import List, Dict, Optional
import time
from PIL import Image
from io import BytesIO
import base64

# 本地Stable Diffusion处理
from diffusers import StableDiffusionPipeline
import torch

class DeepseekService:
    """处理DeepSeek API调用，生成文本分镜描述"""
    
    def __init__(self, api_key=None):
        self.api_key = api_key or os.environ.get("DEEPSEEK_API_KEY")
        self.api_url = "https://api.deepseek.com/v1/chat/completions"
        
    def generate_storyboard_descriptions(self, lyrics: str, style: str = "电影场景", 
                                     segments: int = 5, enhanced_prompt=None) -> List[Dict]:
        """
        基于歌词生成分镜文本描述
        
        Args:
            lyrics: 歌词内容
            style: 分镜风格
            segments: 生成的分镜数量
            enhanced_prompt: 可选的增强提示，用于实现叙事连贯性
            
        Returns:
            分镜描述列表
        """
        if not self.api_key:
            raise ValueError("未设置DeepSeek API密钥")
            
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }
        
        # 构建提示词
        prompt = enhanced_prompt if enhanced_prompt else f"""
    你是一位专业的歌词视觉化专家，请将以下歌词分成{segments}个关键场景，并为每个场景创建详细的视觉描述。
        
        
        
        请根据歌词的情感、故事和意象，以{style}风格创建具有连续性的视觉分镜。
        每个分镜需要包含:
        1. 对应的歌词片段
        2. 详细的场景描述，包括人物、场景、动作、色彩、情感等细节
        3. 场景的视觉氛围和情感基调
        
        请按照以下JSON格式输出:
        ```json
        [
          {{
            "lyrics_segment": "对应的歌词片段",
            "description": "详细的场景描述",
            "mood": "情感基调"
          }},
          ...
        ]
        ```
        只返回JSON格式内容，不要有其他解释。
        """
        
        data = {
            "model": "deepseek-chat",
            "messages": [
                {"role": "system", "content": "你是一位专业的音乐视频分镜导演"},
                {"role": "user", "content": prompt}
            ],
            "temperature": 0.7
        }
        
        try:
            response = requests.post(self.api_url, headers=headers, json=data)
            response.raise_for_status()
            
            result = response.json()
            content = result["choices"][0]["message"]["content"]
            
            # 提取JSON部分
            json_str = content
            if "```json" in content:
                start = content.find("```json") + 7
                end = content.rfind("```")
                json_str = content[start:end].strip()
            
            # 解析JSON
            storyboards = json.loads(json_str)
            return storyboards
            
        except Exception as e:
            logging.error(f"调用DeepSeek API失败: {str(e)}")
            raise

class StableDiffusionService:
    """处理图像生成，基于分镜描述，调用 Stability AI 的 API"""
    
    def __init__(self,api_key, output_dir="app/static/images/storyboards"):
        self.output_dir = output_dir
        self.api_key = api_key or os.environ.get("STABILITY_API_KEY")
        self.api_host = "https://api.stability.ai"
        if not self.api_key:
            raise ValueError("请设置 STABILITY_API_KEY 环境变量")
        
    def generate_image(self, description: str, style: str = "") -> str: 
        """
        基于描述生成图像，调用 Stability AI v2beta 接口
        
        Args:
            description: 场景描述
            style: 额外的风格提示
            
        Returns:
            生成的图像相对于 app/static 文件夹的路径
        """
        # 构建完整的提示词
        prompt = f"{description}. {style} style."
        
        # 使用 v2beta 接口 URL
        endpoint = f"{self.api_host}/v2beta/stable-image/generate/ultra"
        headers = {
            "authorization": f"Bearer {self.api_key}",
            "accept": "image/*"
        }
        # 按照官网示例，需要使用 multipart 方式提交，files 参数提供占位符
        data = {
            "prompt": prompt,
            "output_format": "webp",
        }
        
        response = requests.post(endpoint, headers=headers, files={"none": ""}, data=data, verify=False)
        if response.status_code != 200:
            raise Exception(f"Stability API error: {response.status_code} {response.text}")
        
        os.makedirs(self.output_dir, exist_ok=True)
        image_filename = f"storyboard_{uuid.uuid4()}.webp"
        image_path = os.path.join(self.output_dir, image_filename)
        with open(image_path, "wb") as f:
            f.write(response.content)
        
        # 返回相对于 app/static 的路径，例如 "images/storyboards/xxx.webp"
        relative_path = "images/storyboards/" + image_filename
        return relative_path